[Bugfix][Model] DeepseekV4: resolve kv_cache_dtype="auto" to fp8_ds_mla - #55134
JaredforReal wants to merge 1 commit into
Conversation
The fp8_ds_mla layout requires fp8 KV storage (the FlashMLA decode kernel hardcodes is_fp8_kvcache), yet the DSv4 backends declare "auto" as supported, so starting DeepseekV4 with default options crashed at layer construction: AssertionError: DeepseekV4 fp8_ds_mla layout only supports fp8 kv-cache, got auto Resolve "auto" to fp8_ds_mla in _resolve_dsv4_kv_cache_dtype, writing it back to cache_config so page-size specs pick the 576B slot, and accept "auto" in the SM120 backend validator. This mirrors _canonicalize_sparse_mla_kv_cache_dtype for FLASHMLA_SPARSE / FLASHINFER_MLA_SPARSE_SM120. Explicit non-fp8 dtypes are still rejected. Fixes vllm-project#47174. Un-breaks test_can_initialize_large_subset[DeepseekV4ForConditionalGeneration]. Co-authored-by: Kimi <noreply@moonshot.cn> Signed-off-by: JaredforReal <w13431838023@gmail.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
🟢 Approval recommended
The change is narrowly scoped, aligns validator behavior with the resolver’s canonicalization, and is covered by targeted unit tests.
Pull request overview
Fixes a DeepSeek-V4 initialization crash on backends that require the fp8_ds_mla KV-cache layout by ensuring kv_cache_dtype="auto" is treated as a valid input and is canonicalized to fp8_ds_mla during layer construction (with cache_config write-back for correct page sizing).
Changes:
- Extend
_resolve_dsv4_kv_cache_dtypeto accept"auto"when thefp8_ds_mlalayout is required and canonicalize it to"fp8_ds_mla"(includingcache_config.cache_dtypewrite-back). - Update the SM120 FlashInfer sparse-MLA backend validator to accept
"auto"as an allowedkv_cache_dtype. - Add unit tests covering
"auto"/FP8 alias resolution to"fp8_ds_mla", write-back behavior, and rejection of explicit non-FP8 dtypes under thefp8_ds_mlalayout.
File summaries
| File | Description |
|---|---|
vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py |
Allows "auto" for SM120 validation, relying on later canonicalization to fp8_ds_mla. |
vllm/models/deepseek_v4/attention.py |
Canonicalizes "auto" (and FP8 aliases) to fp8_ds_mla for the packed layout and writes back to cache_config. |
tests/models/test_deepseek_v4_kv_cache_dtype.py |
Adds focused tests for dtype resolution, write-back, and rejection behavior. |
Review details
- Files reviewed: 3/3 changed files
- Comments generated: 0
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
|
Closing as duplicate of #45091 (earlier PR fixing the same For reference, this PR additionally relaxed the SM120 |
Purpose
Fixes #47174.
Starting DeepseekV4 (e.g.
deepseek-ai/DeepSeek-V4-Flash) with default options crashes at layer construction on platforms using thefp8_ds_mlaKV layout (FlashMLA on SM90, FlashInfer sparse on SM120):The DSv4 backends declare
"auto"insupported_kv_cache_dtypes, but nothing resolves it beforeDeepseekV4Attention.__init__asserts on it. The layout architecturally requires fp8 storage (the FlashMLA decode kernel hardcodesis_fp8_kvcache=True), so there is no bf16 fallback for "auto" to resolve to — the only sensible resolution isfp8_ds_mla, which is also DeepSeek's native serving format for this model (the checkpoints are fp8 and all reference configs pass--kv-cache-dtype fp8).This currently breaks
test_can_initialize_large_subset[DeepseekV4ForConditionalGeneration]on the (H200 MIG 35GB) Basic Models (Extra Initialization) shard — reproducible on main, e.g. build 87059.Changes:
_resolve_dsv4_kv_cache_dtype: accept"auto"and resolve it tofp8_ds_mla, writing it back tocache_configso page-size specs pick the 576B per-token slot (same path as the existing"fp8"alias).supports_combinationvalidator: accept"auto"(resolved tofp8_ds_mlaat layer construction).bfloat16) are still rejected, since the kernel has no such implementation.This mirrors
_canonicalize_sparse_mla_kv_cache_dtype, which already resolves"auto"→fp8_ds_mlaforFLASHINFER_MLA_SPARSE_SM120(V3.2 family).Test Plan
New
tests/models/test_deepseek_v4_kv_cache_dtype.py: resolution ofauto/fp8 aliases tofp8_ds_mla,cache_configwrite-back, rejection of explicit non-fp8 dtypes, and plain-layout passthrough.Test Result
No model-eval impact: behavior only changes for
kv_cache_dtype="auto", which previously crashed; explicitfp8*settings resolve exactly as before.Duplicate-work check
Searched open PRs/issues for
fp8_ds_mla auto/_resolve_dsv4_kv_cache_dtype: only issues (#47174, #52938, #47266), no open fix PR found at the time of writing.AI assistance
Prepared with AI assistance (Kimi Code); the human submitter has reviewed every changed line.